from datetime import datetime
from typing import Dict, List
from mcp.server.fastmcp import FastMCP
from wiki_api import WikiSearchAPI

# 创建MCP服务器实例
wiki_mcp = FastMCP("星穹智库服务器")

# 初始化Wiki搜索API
wiki_api = WikiSearchAPI()

# 添加Wiki搜索资源
@wiki_mcp.tool()
def search_wiki(
    query: str, 
    top_k: int = 5,
    min_score: float = 0.3,
    save_results: bool = False
) -> Dict:
    """搜索星穹铁道Wiki内容资源
    
    Args:
        query: 搜索查询文本
        top_k: 返回结果数量(默认5, 范围1-20)
        min_score: 最低相关性分数(默认0.3, 范围0-1)
        save_results: 是否保存搜索结果到文件(默认False)
    
    Returns:
        包含搜索结果和状态的字典
    """
    # 执行搜索
    success, message, results, timestamp = wiki_api.search(
        query=query,
        top_k=top_k,
        min_score=min_score
    )
    
    # 可选地保存结果
    if success and save_results:
        filename = f"wiki_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        wiki_api.save_search_results(results, filename)
    
    return {
        "status": "success" if success else "error",
        "message": message,
        "query": query,
        "count": len(results),
        "results": results,
        "timestamp": timestamp,
        "metadata": {
            "top_k": top_k,
            "min_score": min_score,
            "results_saved": save_results and success
        }
    }


def run():
    """运行MCP服务器"""
    wiki_mcp.run(transport="sse")

if __name__ == "__main__":
    run()
